import shutil
import xml.etree.ElementTree as ET
import os
import random


def get_all_classes(xml_dir):
    class_list = []
    xml_list = os.listdir(xml_dir)
    for xml in xml_list:
        xml_file = os.path.join(xml_dir, xml)
        file_tree = ET.parse(xml_file)
        tree_root = file_tree.getroot()
        object_element = tree_root.findall("object")
        for obj in object_element:
            if obj.find('name').text not in class_list:
                class_list.append(obj[0].text)
    return class_list


def convert_format(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    if w >= 1:
        w = 0.99
    if h >= 1:
        h = 0.99
    return x, y, w, h


def convert_annotation(xml_dir, xml_name, txt_dir, classes):
    xml_file = os.path.join(xml_dir, xml_name)
    with open(xml_file, "r", encoding='UTF-8') as in_file:
        txt_name = xml_name[:-4] + '.txt'
        print(txt_name)
        # txt_storage = os.path.join(storage_dir, 'JPEGImages/labels')
        txt_file = os.path.join(txt_dir, txt_name)
        with open(txt_file, "w+", encoding='UTF-8') as out_file:
            tree = ET.parse(in_file)
            root = tree.getroot()
            size = root.find('size')
            w = int(size.find('width').text)
            h = int(size.find('height').text)
            out_file.truncate()
            for obj in root.iter('object'):
                difficult = obj.find('difficult').text
                cls = obj.find('name').text
                if cls not in classes or int(difficult) == 1:
                    continue
                cls_id = classes.index(cls)
                xml_box = obj.find('bndbox')
                b = (
                    float(xml_box.find('xmin').text), float(xml_box.find('xmax').text),
                    float(xml_box.find('ymin').text),
                    float(xml_box.find('ymax').text))
                bb = convert_format((w, h), b)
                out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


if __name__ == "__main__":
    # classes = ['baggage', 'baggage-on', 'guardian', 'guest', 'head']
    voc_dir = r"/home/kaijia/algo-env/datasets/hq_safeguard_dataset/VOC"
    # classes = get_all_classes(r"/home/kaijia/algo-env/datasets/hq_safeguard_dataset/VOC/Annotations")
    classes = ['aeroplane', 'airplane-head', 'bridge-head', 'cabin-door',
               'car_ladder', 'car_refuel', 'car_stars', 'cargo-door',
               'chock', 'connecting', 'dinning_car', 'guide_car',
               'person', 'tire', 'tractor_car', 'transfer_car']
    yolo_dir = r"/home/kaijia/algo-env/datasets/hq_safeguard_dataset/yolov7"
    if not os.path.exists(yolo_dir):
        os.makedirs(yolo_dir)
    img_voc = os.path.join(voc_dir, 'Images')
    img_yolo = os.path.join(yolo_dir, 'Images')
    shutil.copytree(img_voc, img_yolo)
    xml_dir = os.path.join(voc_dir, "Annotations")
    txt_dir = os.path.join(yolo_dir, "Labels")
    if not os.path.exists(txt_dir):
        os.makedirs(txt_dir)
    xml_list = os.listdir(xml_dir)
    for i in range(0, len(xml_list)):
        file_full_path = os.path.join(xml_dir, xml_list[i])
        if ('.xml' in file_full_path) or ('.XML' in file_full_path):
            convert_annotation(xml_dir, xml_list[i], txt_dir, classes)
        else:
            print('not xml file', i)
